{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 315
    },
    "colab_type": "code",
    "id": "c4hVOeZ69doH",
    "outputId": "1773473b-cc1c-417b-fab4-86c724345824"
   },
   "outputs": [],
   "source": [
    "# load pretrained model weights\n",
    "!wget https://github.com/wielandbrendel/robustness_workshop/releases/download/v0.0.1/diversity_cifar10_ResNet20v1_model.159.h5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "_Vhv1ef4EG6A",
    "outputId": "95aa364f-780f-42be-e232-1eea421b748e"
   },
   "outputs": [],
   "source": [
    "# make sure the right version of tensorflow is installed\n",
    "!pip3 install --upgrade tensorflow==2.0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "sI2YD4T-9IBf"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "from tensorflow.keras.layers import AveragePooling2D, Input, Flatten\n",
    "from tensorflow.keras.models import Model, load_model\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "import foolbox as fb\n",
    "from model import resnet_v1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9iWEGTbM9IBs"
   },
   "outputs": [],
   "source": [
    "# parameters\n",
    "num_classes = 10\n",
    "lamda = 2.0\n",
    "log_det_lamda = 0.5\n",
    "augmentation = False\n",
    "filepath = 'diversity_cifar10_ResNet20v1_model.159.h5'\n",
    "\n",
    "n = 3\n",
    "depth = n * 6 + 2\n",
    "version = 1\n",
    "\n",
    "input_shape = (32, 32, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 54
    },
    "colab_type": "code",
    "id": "mqVbMl0G9IB2",
    "outputId": "f01231bf-535c-4435-b3e5-7357fae5ab29",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model_input = Input(shape=input_shape)\n",
    "model_dic = {}\n",
    "model_out = []\n",
    "for i in range(3):\n",
    "    model_dic[str(i)] = resnet_v1(input=model_input, depth=depth, num_classes=num_classes, dataset='cifar10')\n",
    "    model_out.append(model_dic[str(i)][2])\n",
    "model_output = tf.keras.layers.concatenate(model_out)\n",
    "model = Model(inputs=model_input, outputs=model_output)\n",
    "model_ensemble = tf.keras.layers.Average()(model_out)\n",
    "model_ensemble = Model(inputs=model_input, outputs=model_ensemble)\n",
    "\n",
    "# load model\n",
    "model.load_weights(filepath)\n",
    "\n",
    "# compile model\n",
    "model_ensemble.compile('sgd')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 69
    },
    "colab_type": "code",
    "id": "Rp-ylXuQ9ICA",
    "outputId": "2ef8745c-1174-4491-fa3b-cfd692a4f850"
   },
   "outputs": [],
   "source": [
    "# Subtracting pixel mean improves accuracy\n",
    "subtract_pixel_mean = True\n",
    "\n",
    "import keras\n",
    "from keras.datasets import cifar10\n",
    "\n",
    "# Load the data.\n",
    "(x_train, y_train), (x_test, y_test) = cifar10.load_data()\n",
    "\n",
    "# Input image dimensions.\n",
    "input_shape = x_train.shape[1:]\n",
    "\n",
    "# Normalize data.\n",
    "x_train = x_train.astype('float32') / 255\n",
    "x_test = x_test[:200].astype('float32') / 255\n",
    "\n",
    "# If subtract pixel mean is enabled\n",
    "clip_min = 0.0\n",
    "clip_max = 1.0\n",
    "if subtract_pixel_mean:\n",
    "    x_train_mean = np.mean(x_train, axis=0)\n",
    "    x_train -= x_train_mean\n",
    "    x_test -= x_train_mean\n",
    "    clip_min -= x_train_mean\n",
    "    clip_max -= x_train_mean\n",
    "    \n",
    "y_test, y_train = y_test[:200].flatten(), y_train.flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "WQ___vx59ICG",
    "outputId": "c895e3f1-873f-479f-c749-f68470ba2c62"
   },
   "outputs": [],
   "source": [
    "pred = model_ensemble.predict(x_test).argmax(1)\n",
    "print(f'Clean accuracy: {np.mean(pred == y_test):.3f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "OUvwu5zz9ICU"
   },
   "outputs": [],
   "source": [
    "# convert to Foolbox model\n",
    "fmodel = fb.models.TensorFlowModel(model_ensemble, bounds=(-2, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "GnhjZMJ89ICY"
   },
   "outputs": [],
   "source": [
    "images = tf.convert_to_tensor(x_test, dtype=tf.float32)\n",
    "labels = tf.convert_to_tensor(y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# a simple wrapper for the init attack in BB\n",
    "class init_attack(object):\n",
    "    \n",
    "    def __init__(self, attack):\n",
    "        self.attack = attack\n",
    "        \n",
    "    def run(self, model, originals, criterion_):\n",
    "        return self.attack(model, images, criterion=criterion_, epsilons=0.3)[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "acc = 0\n",
    "total_images = 0\n",
    "\n",
    "epsilon = 0.01\n",
    "repetitions = 3\n",
    "\n",
    "pdg20_attack = fb.attacks.LinfPGD(steps=20, abs_stepsize=epsilon/10, random_start=True)\n",
    "pdg200_attack = fb.attacks.LinfPGD(steps=200, abs_stepsize=epsilon/20, random_start=True)\n",
    "pdg_init_attack = fb.attacks.LinfPGD(steps=20, abs_stepsize=epsilon/2, random_start=True)\n",
    "bb_attack = fb.attacks.LinfinityBrendelBethgeAttack(init_attack(pdg_init_attack), steps=200)\n",
    "\n",
    "for _images, _labels in zip(np.split(x_test, 10), np.split(y_test, 10)):\n",
    "    mask = np.array([True] * 20)\n",
    "    images = tf.convert_to_tensor(_images[mask], dtype=tf.float32)\n",
    "    labels = tf.convert_to_tensor(_labels[mask])\n",
    "    \n",
    "    _epsilon = 10 * epsilon\n",
    "    \n",
    "    # run quick PGD attacks\n",
    "    for r in range(repetitions):\n",
    "        if mask.sum() > 0:\n",
    "            adv, adv_clipped, adv_mask = pdg20_attack(fmodel, images, criterion=fb.criteria.Misclassification(labels), epsilons=epsilon)\n",
    "        \n",
    "            mask[mask] = ~adv_mask.numpy()\n",
    "\n",
    "            images = tf.convert_to_tensor(_images[mask], dtype=tf.float32)\n",
    "            labels = tf.convert_to_tensor(_labels[mask])\n",
    "        \n",
    "        \n",
    "    # run longer PGD attacks\n",
    "    for r in range(repetitions):\n",
    "        if mask.sum() > 0:\n",
    "            adv, adv_clipped, adv_mask = pdg200_attack(fmodel, images, criterion=fb.criteria.Misclassification(labels), epsilons=epsilon)\n",
    "        \n",
    "            mask[mask] = ~adv_mask.numpy()\n",
    "\n",
    "            images = tf.convert_to_tensor(_images[mask], dtype=tf.float32)\n",
    "            labels = tf.convert_to_tensor(_labels[mask])\n",
    "    \n",
    "        \n",
    "    # run Brendel & Bethge attack\n",
    "    for r in range(repetitions):\n",
    "        if mask.sum() > 0:\n",
    "            adv, adv_clipped, adv_mask = bb_attack(fmodel, images, criterion=fb.criteria.Misclassification(labels), epsilons=epsilon)\n",
    "        \n",
    "            mask[mask] = ~adv_mask.numpy()\n",
    "\n",
    "            images = tf.convert_to_tensor(_images[mask], dtype=tf.float32)\n",
    "            labels = tf.convert_to_tensor(_labels[mask])\n",
    "    \n",
    "    acc += (1 - adv_mask.numpy().mean()) * len(adv)\n",
    "    total_images += _images.shape[0]\n",
    "    \n",
    "    print(total_images, acc / total_images)\n",
    "\n",
    "print(f'Model accuracy on adversarials: {acc / total_images:.3f}')"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "diversity_attack.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
